from utils import *
import os
import sys
import numpy as np
from threading import Lock
import queue
from models.base_model import BaseModel, ModelConfig, LstmNet
import torch
import numpy as np
from config import config_model_singleton
from utils import sys_logger

# sys.path.append(os.path.abspath(os.path.join(os.path.dirname("__file__"), '..')))


class Base_Servo_AnomalyDetector(BaseModel):
    """
    伺服类产品的基础类
    """

    def __init__(self, config: ModelConfig):
        super().__init__(config)

        # 输入变量相关
        self.input_dim = 0
        self.input_update_flag = dict()
        for k, v in config.inputs.items():
            # (flag, offset)
            self.input_update_flag[k] = [False, self.input_dim]
            self.input_dim += len(v)

        # mean and std
        self.mean = np.array(config.mean)
        self.std = np.array(config.std)
        self.need_preprocess = self._config.inputs_need_preprocess

        # dim0 为batchsize
        self.input_data = np.zeros((1, self.input_dim))  # 初始化为零

        self.output_data_queue = queue.Queue(maxsize=config.input_output_queue_size)
        self.input_data_queue = queue.Queue(maxsize=config.input_output_queue_size)
        # 检查cuda是否可用
        self.device = torch.device("cuda:0" if config_singleton.use_cuda and torch.cuda.is_available() else "cpu")

        self.hidden_pred = None

        # 线程锁
        self.lock = Lock()
        self.thread_stop = False

    def update_input_data(self, msgbean: MsgBean):
        """
        更新输入数据.
        根据数据所在的位置和所在的表号，更新输入数据.
        """
        table_num = msgbean.table_num
        # 获取当前表的偏移量, 存储在inputdata中的偏移量
        offset = self.input_update_flag[table_num][1]

        # cnt = len(self._config.inputs[table_num])
        with self.lock:
            # 根据数组中的值作为字典的key, 取出对应key的value
            # numpy 可以使用切片的方式。
            # self.input_data[0][offset: offset +
            #                    cnt] = msgbean.vars[self._config.inputs[table_num]]
            for index in self._config.inputs[table_num]:
                self.input_data[0][offset] = msgbean.vars[index]
                offset += 1
            # 更新标志位
            self.input_update_flag[table_num][0] = True

        if all(flag for [flag, _] in self.input_update_flag.values()):
            # TODO 考虑是否使用带时间，推理结束后无法对应于那个时间的数据
            self.input_data_queue.put([msgbean.vars[1], self.input_data.copy()])
            # 数据放入了队列中，便可以把标志清理掉。
            # 先将标志重置，再进行推理
            with self.lock:
                for k, v in self.input_update_flag.items():
                    v[0] = False

    def train(self, data):
        return super().train(data)

    def run(self):
        """
        模型推理线程.
        一帧数据进来，就进行模型推理，并将推理结果放入输出队列中.
        需要启动多线程，一个线程负责接收数据。
        """
        while not self.thread_stop:
            # 阻塞式等待数据的到来。
            t, input = self.input_data_queue.get()

            if input is None or t is None:
                continue

            # 输入数据归一化转换
            if self.need_preprocess:
                input = (input - self.mean) / self.std
            # 一步推理过程
            test_x = torch.from_numpy(input).float().to(self.device)
            pred_x, self.hidden_pred = self.model(test_x, self.hidden_pred)
            # convert pred_x to numpy
            pred_x = pred_x.detach().cpu().numpy()
            # 逆归一化, 输出的数据比输入的数据少一维
            if self.need_preprocess:
                pred_x = pred_x * self.std[1:] + self.mean[1:]

            # TODO:
            # 增加异常检测功能

            # 将时间一起放进队列中
            self.output_data_queue.put([t, pred_x[0]])

        sys_logger.info(f"{self._config.name} Model inference thread stopped.")

    def load_weights(self, weights_path):
        """
        load pytorch pth file to model
        """
        real_path = os.path.join(os.path.dirname(
            __file__), '..', 'model_files', weights_path)
        if not os.path.exists(real_path):
            raise FileNotFoundError(
                "Hydraulic Servo Model file not find, check file path or name.")

        if self.model.__class__.__name__ == 'LstmNet':
            # self.model = load(real_path)
            self.model.load_state_dict(torch.load(
                real_path, map_location=self.device))

    def stop(self):
        self.thread_stop = True
        self.input_data_queue.put([None, None])
        self.output_data_queue.put([None, None])

    # 同时获取输出变量的名字和数值
    def get_output_varnames(self):

        return self._config.outputs["var_names"]

    def get_output_data(self):
        """
        获取异常检测的输出结果， 带时间标签，如果没有数据立即返回。
        """
        if self.output_data_queue.empty():
            vals = [None, None]
        else:
            vals = self.output_data_queue.get_nowait()

        return vals

# Example usage
if __name__ == "__main__":

    # config example
    pass
